import numpy as np
import matplotlib.pyplot as plt
from scipy.linalg import expm

# 设置支持中文的字体
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False


class DroneAttitude:
    def __init__(self, dt=0.01):
        # 状态变量: [角度, 角速度]
        self.state = np.array([0.0, 0.0])  # 初始角度和角速度
        self.dt = dt  # 时间步长

        # 物理参数
        self.I = 0.01  # 转动惯量
        self.damping = 0.1  # 阻尼系数

    def update(self, torque):
        # 物理模型更新: 欧拉积分
        # 角加速度 = (扭矩 - 阻尼*角速度) / 转动惯量
        angular_acc = (torque - self.damping * self.state[1]) / self.I

        # 更新状态
        self.state[1] += angular_acc * self.dt  # 更新角速度
        self.state[0] += self.state[1] * self.dt  # 更新角度

        return self.state.copy()


class KalmanFilter:
    def __init__(self, dt, process_noise_std, measurement_noise_std):
        # 状态转移矩阵 (角度, 角速度)
        self.F = np.array([[1, dt],
                           [0, 1]])

        # 控制输入矩阵 (这里假设控制输入直接作用于角速度)
        self.B = np.array([[0],
                           [dt]])

        # 测量矩阵 (假设我们只能直接测量角度)
        self.H = np.array([[1, 0]])

        # 过程噪声协方差
        self.Q = np.diag(process_noise_std) ** 2

        # 测量噪声协方差
        self.R = np.array([[measurement_noise_std ** 2]])

        # 初始状态协方差
        self.P = np.eye(2)

        # 初始状态
        self.x = np.array([[0],  # 角度
                           [0]])  # 角速度

    def predict(self, u=0):
        # 状态预测
        self.x = self.F @ self.x + self.B * u

        # 协方差预测
        self.P = self.F @ self.P @ self.F.T + self.Q

        return self.x

    def update(self, z):
        # 计算卡尔曼增益
        S = self.H @ self.P @ self.H.T + self.R
        K = self.P @ self.H.T @ np.linalg.inv(S)

        # 状态更新
        y = z - self.H @ self.x
        self.x = self.x + K @ y

        # 协方差更新
        I = np.eye(self.P.shape[0])
        self.P = (I - K @ self.H) @ self.P

        return self.x


# 模拟参数
dt = 0.01  # 时间步长
total_time = 10  # 总模拟时间
steps = int(total_time / dt)  # 总步数

# 创建无人机模型
drone = DroneAttitude(dt)

# 创建卡尔曼滤波器
# 过程噪声标准差: [角度噪声, 角速度噪声]
# 测量噪声标准差: 角度测量噪声
kf = KalmanFilter(dt, process_noise_std=[0.01, 0.1], measurement_noise_std=0.1)

# 存储数据
true_angles = []
true_angular_velocities = []
gyro_measurements = []
accel_measurements = []
kf_estimates = []
kf_velocity_estimates = []
time_points = []

# 模拟循环
for t in range(steps):
    current_time = t * dt
    time_points.append(current_time)

    # 生成控制扭矩 (正弦波)
    torque = 0.1 * np.sin(2 * current_time)

    # 更新无人机真实状态
    true_state = drone.update(torque)
    true_angles.append(true_state[0])
    true_angular_velocities.append(true_state[1])

    # 模拟传感器测量 (添加噪声)
    gyro_noise = np.random.normal(0, 0.05)  # 陀螺仪噪声
    accel_noise = np.random.normal(0, 0.1)  # 加速度计噪声

    # 陀螺仪测量角速度
    gyro_measurement = true_state[1] + gyro_noise
    gyro_measurements.append(gyro_measurement)

    # 加速度计测量角度 (通过重力分量)
    accel_measurement = true_state[0] + accel_noise
    accel_measurements.append(accel_measurement)

    # 卡尔曼滤波预测步骤 (使用陀螺仪测量作为控制输入)
    kf.predict(u=gyro_measurement)

    # 卡尔曼滤波更新步骤 (使用加速度计测量)
    kf_estimate = kf.update(np.array([[accel_measurement]]))
    kf_estimates.append(kf_estimate[0, 0])
    kf_velocity_estimates.append(kf_estimate[1, 0])

# 绘制结果
plt.figure(figsize=(12, 10))

# 角度比较
plt.subplot(2, 1, 1)
plt.plot(time_points, true_angles, label='真实角度', linewidth=2)
plt.plot(time_points, accel_measurements, label='加速度计测量', alpha=0.7)
plt.plot(time_points, kf_estimates, label='卡尔曼滤波估计', linewidth=2)
plt.xlabel('时间 (s)')
plt.ylabel('角度 (rad)')
plt.title('无人机俯仰角估计')
plt.legend()
plt.grid(True)

# 角速度比较
plt.subplot(2, 1, 2)
plt.plot(time_points, true_angular_velocities, label='真实角速度', linewidth=2)
plt.plot(time_points, gyro_measurements, label='陀螺仪测量', alpha=0.7)
plt.plot(time_points, kf_velocity_estimates, label='卡尔曼滤波估计', linewidth=2)
plt.xlabel('时间 (s)')
plt.ylabel('角速度 (rad/s)')
plt.title('无人机角速度估计')
plt.legend()
plt.grid(True)

plt.tight_layout()
plt.show()

# 计算并显示误差
accel_errors = np.array(accel_measurements) - np.array(true_angles)
gyro_errors = np.array(gyro_measurements) - np.array(true_angular_velocities)
kf_angle_errors = np.array(kf_estimates) - np.array(true_angles)
kf_velocity_errors = np.array(kf_velocity_estimates) - np.array(true_angular_velocities)

print(f"加速度计角度误差标准差: {np.std(accel_errors):.4f} rad")
print(f"陀螺仪角速度误差标准差: {np.std(gyro_errors):.4f} rad/s")
print(f"卡尔曼滤波角度误差标准差: {np.std(kf_angle_errors):.4f} rad")
print(f"卡尔曼滤波角速度误差标准差: {np.std(kf_velocity_errors):.4f} rad/s")